package handler

import (
	"context"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
	"myai/core"
	"myai/core/types"
	"myai/service/openai/chat"
	"myai/store/model"
	"myai/store/model/req"
	"myai/store/vo"
	"myai/utils"
	"myai/utils/resp"
)

type ChatOpenAiHandler struct {
	BaseHandler
	openAiChatService *chat.OpenAiChatService
}

func NewChatOpenAiHandler(app *core.AppServer, db *gorm.DB, openAiChatService *chat.OpenAiChatService) *ChatOpenAiHandler {
	return &ChatOpenAiHandler{
		BaseHandler:       BaseHandler{App: app, DB: db},
		openAiChatService: openAiChatService,
	}
}

// Tokens godoc
// @Summary      统计 token 数量
// @Description  统计 token 数量
// @Tags         前台-用户会话
// @Accept       json
// @Produce      json
// @Param        req   body     req.ChatTokens  true  "统计tokens数量参数"
// @Success      200  {object}  types.BizVo{data=int}
// @Router       /api/chat/tokens [post]
// @Security ApiKeyAuth
func (h *ChatOpenAiHandler) Tokens(c *gin.Context) {
	var data req.ChatTokens
	if err := c.ShouldBindJSON(&data); err != nil {
		resp.ERROR(c, types.InvalidArgs)
		return
	}
	// 计算token
	tokens, err := utils.CalcTokens(data.Text, data.Model)
	if err != nil {
		log.Errorf("计算token失败：%v", err)
		return
	}
	resp.SUCCESS(c, tokens)
}

// GetMsgTokens godoc
// @Summary      根据消息Id查询消息信息
// @Description  根据消息Id查询消息信息
// @Tags         前台-用户会话
// @Accept       json
// @Produce      json
// @Param        req   body     req.ChatGetMsgTokens  true  "获取消息详情"
// @Success      200  {object}  types.BizVo{data=[]vo.HistoryMessage}
// @Router       /api/chat/getMsgTokens [post]
// @Security ApiKeyAuth
func (h *ChatOpenAiHandler) GetMsgTokens(c *gin.Context) {
	var data req.ChatGetMsgTokens
	if err := c.ShouldBindJSON(&data); err != nil {
		resp.ERROR(c, types.InvalidArgs)
		return
	}

	var messagesDB []model.ChatMessage
	var messageVos = make([]vo.HistoryMessage, 0)
	// 查询消息
	err := h.DB.Where("id IN ?", data.MsgIds).Find(&messagesDB).Error
	if err != nil {
		resp.ERROR(c, "操作数据库出现错误："+err.Error())
		return
	} else {
		for _, item := range messagesDB {
			var msgVo vo.HistoryMessage
			err = utils.CopyObject(item, &msgVo)
			if err != nil {
				log.Errorf("拷贝数据出现错误：%v", err)
				continue
			}
			msgVo.Id = item.Id
			msgVo.CreatedAt = item.CreatedAt.Unix()
			msgVo.UpdatedAt = item.CreatedAt.Unix()
			messageVos = append(messageVos, msgVo)
		}
	}
	resp.SUCCESS(c, messageVos)
}

// StopGenerate godoc
// @Summary      停止生成
// @Description  停止生成
// @Tags         前台-用户会话
// @Accept       json
// @Produce      json
// @Param        session_id   query     string  true  "会话Id"
// @Success      200  {object}  types.BizVo{}
// @Router       /api/chat/stop [get]
// @Security ApiKeyAuth
func (h *ChatOpenAiHandler) StopGenerate(c *gin.Context) {
	sessionId := h.GetTrim(c, "session_id")
	/*if h.ReqCancelFunc.Has(sessionId) {
		h.ReqCancelFunc.Get(sessionId)()
		h.ReqCancelFunc.Delete(sessionId)
	}*/
	h.openAiChatService.StopGenerate(sessionId)
	resp.SUCCESS(c, types.OkMsg)
}

func (h *ChatOpenAiHandler) Chat(ctx context.Context, session *types.ChatSession, role model.ChatRole, prompt string, ws *types.WsClient) error {
	return h.openAiChatService.Chat(ctx, session, role, prompt, ws)
}

func (h *ChatOpenAiHandler) SetCtx(clientId string, cancel context.CancelFunc) {
	h.openAiChatService.ReqCancelFunc.Put(clientId, cancel)
}

func (h *ChatOpenAiHandler) ClearChatCtx(chatId string) {
	h.openAiChatService.ChatContexts.Delete(chatId)
}
